import argparse
import cv2
import numpy as np
import os

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer
from basicsr.utils import imwrite

def get_instance():
  """
  初始化模型
  """
  restorer_model_name = 'GFPGANv1.4'
  url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'
  arch = 'clean'
  channel_multiplier = 2

  # determine model paths
  model_path = os.path.join('models', restorer_model_name + '.pth')
  if not os.path.isfile(model_path):
      # download pre-trained models from url
      model_path = url

  bg_upsampler = None
  restorer = GFPGANer(
      model_path=model_path,
      upscale=2,
      arch=arch,
      channel_multiplier=channel_multiplier,
      bg_upsampler=bg_upsampler)
  
  return restorer
#end def

def image_restorer(image_path):
  """
  图像高清修复函数
  """
  img_name = os.path.basename(image_path)
  print(f'Processing {img_name} ...')
  basename, ext = os.path.splitext(img_name)
  input_img = cv2.imread(image_path, cv2.IMREAD_COLOR)

  restorer = get_instance()
  cropped_faces, restored_faces, restored_img = restorer.enhance(
    input_img,
    has_aligned=False,
    only_center_face=False,
    paste_back=True,
    weight=0.5
  )

  output = 'output_restorer'
  extension = 'png'
  # save restored img
  if restored_img is not None:
    save_restore_path = os.path.join(output, f'{basename}.{extension}')
    imwrite(restored_img, save_restore_path)
#end def